from volcenginesdkarkruntime import Ark
from openai import OpenAI
from pathlib import Path
from log import log



class DouBao:
    def __init__(self):
        self.client = Ark(
        sk="AKLTOThkMDRiZWI2ZTNlNGY1YmE0NzRmNDE5OTQyYmU3ZWU",
        ak="WTJOaU5UWTFOV1V5TkRsaE5ETTNaR0ZsTURFNE9EazFNVGxsTm1JeVlqVQ==",
        api_key='78018cba-25ba-4e0b-ac6d-5ab1a2abecf6'
    )


    def doubao_chat(self,question) -> str:
        completion = self.client.chat.completions.create(
            model="ep-20250219142727-tspkh",
            messages=[
                {"role": "system", "content": "你是豆包，是由字节跳动开发的 AI 人工智能助手"},
                {"role": "user", "content": question},
            ],
        )
        return completion.choices[0].message.content

class Kimi:
    def __init__(self):
        self.client = OpenAI(
            api_key="sk-HrSOhl9oIvJKZqcdiqGaYqu6tXMXzA28sJeRRV7Wrejt3WqG",
            base_url="https://api.moonshot.cn/v1",
        )

    def kimi_chat(self,question) -> str:
        completion = self.client.chat.completions.create(
            model="moonshot-v1-8k",
            messages=[
                {"role": "system",
                 "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。"},
                {"role": "user", "content": question},
            ],
            temperature=0.3,
        )
        return completion.choices[0].message.content

    def kimi_file(self,file_path, question):
        log.info("进入kimi抽取文件阶段")
        file_object = self.client.files.create(file=Path(file_path), purpose="file-extract")
        file_id = file_object.id
        file_content = self.client.files.content(file_id=file_id).text
        messages = [
            {
                "role": "system",
                "content": "你是 Kimi，由 Moonshot AI 提供的人工智能助手，你更擅长中文和英文的对话。你会为用户提供安全，有帮助，准确的回答。同时，你会拒绝一切涉及恐怖主义，种族歧视，黄色暴力等问题的回答。Moonshot AI 为专有名词，不可翻译成其他语言。",
            },
            {
                "role": "system",
                "content": file_content,
            },
            {"role": "user", "content": question},
        ]

        completion = self.client.chat.completions.create(
            model="moonshot-v1-32k",
            messages=messages,
            temperature=0.1,
        )
        self.client.files.delete(file_id=file_id)
        return completion.choices[0].message.content
kimi = Kimi()

if __name__ == '__main__':
    kimi = Kimi()
    kimi_chat = kimi.kimi_chat("中国首都？")
    print(kimi_chat)